#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import pytz
import time
from datetime import datetime, timedelta
from typing import Optional

# 设置时区环境变量
os.environ['TZ'] = 'Asia/Shanghai'
if hasattr(time, 'tzset'):
    time.tzset()

# 中国时区常量
CHINA_TZ = pytz.timezone('Asia/Shanghai')


def now_china() -> datetime:
    """
    获取当前中国时间
    
    Returns:
        datetime: 带有中国时区信息的当前时间
    """
    return datetime.now(CHINA_TZ)


def utc_to_china(utc_dt: datetime) -> datetime:
    """
    将UTC时间转换为中国时间
    
    Args:
        utc_dt: UTC时间
        
    Returns:
        datetime: 中国时间
    """
    if utc_dt.tzinfo is None:
        utc_dt = pytz.utc.localize(utc_dt)
    return utc_dt.astimezone(CHINA_TZ)


def china_to_utc(china_dt: datetime) -> datetime:
    """
    将中国时间转换为UTC时间
    
    Args:
        china_dt: 中国时间
        
    Returns:
        datetime: UTC时间
    """
    if china_dt.tzinfo is None:
        china_dt = CHINA_TZ.localize(china_dt)
    return china_dt.astimezone(pytz.utc)


def timestamp_to_china(timestamp: float) -> datetime:
    """
    将时间戳转换为中国时间
    
    Args:
        timestamp: 时间戳（秒）
        
    Returns:
        datetime: 中国时间
    """
    return datetime.fromtimestamp(timestamp, CHINA_TZ)


def timestamp_ms_to_china(timestamp_ms: int) -> datetime:
    """
    将毫秒时间戳转换为中国时间
    
    Args:
        timestamp_ms: 毫秒时间戳
        
    Returns:
        datetime: 中国时间
    """
    return datetime.fromtimestamp(timestamp_ms / 1000, CHINA_TZ)


def china_timestamp() -> int:
    """
    获取当前中国时间的时间戳
    
    Returns:
        int: 时间戳（秒）
    """
    return int(now_china().timestamp())


def localize_datetime(dt: datetime, is_china: bool = True) -> datetime:
    """
    为没有时区信息的datetime对象添加时区信息
    
    Args:
        dt: 没有时区信息的datetime对象
        is_china: 是否为中国时区，默认为True
        
    Returns:
        datetime: 带有时区信息的datetime对象
    """
    if dt.tzinfo is not None:
        return dt
    
    if is_china:
        return CHINA_TZ.localize(dt)
    else:
        return pytz.utc.localize(dt)


def format_china_time(dt: Optional[datetime] = None, format_str: str = '%Y-%m-%d %H:%M:%S') -> str:
    """
    格式化中国时间为字符串
    
    Args:
        dt: datetime对象，如果为None则使用当前时间
        format_str: 格式化字符串
        
    Returns:
        str: 格式化后的时间字符串
    """
    if dt is None:
        dt = now_china()
    elif dt.tzinfo is None:
        dt = CHINA_TZ.localize(dt)
    elif dt.tzinfo != CHINA_TZ:
        dt = dt.astimezone(CHINA_TZ)
    
    return dt.strftime(format_str)


# 兼容性函数，替代原有的时间函数
def get_current_timestamp() -> int:
    """
    获取当前时间戳，使用中国时区
    
    Returns:
        int: 当前时间戳
    """
    return china_timestamp()
